from datetime import datetime

from flask import request

from air_web.web_flask.bll.predict_bll import PredictBll
from air_web.web_flask.common.constants import (
    PAGE_NO,
    PAGE_SIZE,
    fk_predict_columns_chinese_dict,
    predict_columns_chinese_dict,
)
from air_web.web_flask.common.validator import (
    T_CHINESE,
    T_DATE,
    T_INT,
    T_LIST,
    T_STR,
    Fields,
    with_validator,
)
from air_web.web_flask.views.base_page import BasePage


class PredictPageValidationFields:
    input_fields = [
        Fields("area_code", T_STR, must=False),
        Fields("type_id", T_INT, must=False),
        Fields("start_date", T_DATE, must=True),
        Fields("download", T_STR, must=False, default="false"),
    ]
    max_fields = [
        Fields("area_code", T_STR, must=False),
        Fields("start_date", T_DATE, must=True),
        Fields("end_date", T_DATE, must=True),
        Fields("download", T_STR, must=False, default="false"),
    ]


class PredictPage(BasePage):
    """总览页面视图函数类"""

    @classmethod
    @with_validator(request, PredictPageValidationFields.input_fields)
    def industry_predict(cls):
        validation_data = request.validation_data
        is_download, validation_data = BasePage.get_is_download(
            validation_data
        )
        result = PredictBll().get_all_result(validation_data, is_download)

        if is_download:
            if result.empty:
                return cls.get_return_excel_file_from_df(
                    result,
                    {},
                    sheet_name="行业负荷预测数据",
                    file_name="air_predict_table",
                )
            columns = result.columns.to_list()
            two_column_list = list(
                set(columns) - set(list(predict_columns_chinese_dict.keys()))
            )
            times = [
                datetime.strptime(time_str, "%H:%M")
                for time_str in two_column_list
            ]
            sorted_times = sorted(times)
            two_column_list = [time.strftime("%H:%M") for time in sorted_times]
            if two_column_list:
                predict_columns_chinese_dict.update(
                    {col: col for col in two_column_list}
                )
            return cls.get_return_excel_file_from_df(
                result,
                predict_columns_chinese_dict,
                sheet_name="行业负荷预测数据",
                file_name="air_predict_table",
                round_two_decimal_column=two_column_list,
            )
        return cls.return_common_func(result)

    @classmethod
    @with_validator(request, PredictPageValidationFields.input_fields)
    def fk_predict(cls):
        validation_data = request.validation_data
        is_download, validation_data = BasePage.get_is_download(
            validation_data
        )
        result = PredictBll().get_all_fk_result(validation_data, is_download)

        if is_download:
            if result.empty:
                return cls.get_return_excel_file_from_df(
                    result,
                    {},
                    sheet_name="总负荷预测数据",
                    file_name="air_predict_table",
                )

            columns = result.columns.to_list()
            two_column_list = list(
                set(columns)
                - set(list(fk_predict_columns_chinese_dict.keys()))
            )
            times = [
                datetime.strptime(time_str, "%H:%M")
                for time_str in two_column_list
            ]
            sorted_times = sorted(times)
            two_column_list = [time.strftime("%H:%M") for time in sorted_times]
            fk_predict_columns_chinese_dict.update(
                {col: col for col in two_column_list}
            )

            return cls.get_return_excel_file_from_df(
                result,
                fk_predict_columns_chinese_dict,
                sheet_name="总负荷预测数据",
                file_name="air_predict_table",
                round_two_decimal_column=two_column_list,
            )
        return cls.return_common_func(result)

    @classmethod
    @with_validator(request, PredictPageValidationFields.max_fields)
    def max_p_predict(cls):
        validation_data = request.validation_data
        is_download, validation_data = BasePage.get_is_download(
            validation_data
        )
        result = PredictBll().get_max_p_result(validation_data, is_download)
        if is_download:
            columns_chinese_dict = {
                "org_name": "单位名称",
                "data_time": "日期",
                "p_total_pre": "预测最大负荷",
                "p_kt_pre": "预测空调最大负荷",
                "tmp_pre": "预测温度",
                "p_total_true": "最大负荷",
                "p_kt_true": "空调最大负荷",
                "tmp_true": "温度",
                "pre_p_total_rate": "预测最大负荷准确率(%)",
                "pre_p_kt_rate": "预测空调最大负荷准确率(%)",
            }
            return cls.get_return_excel_file_from_df(
                data_df=result,
                columns_chinese_dict=columns_chinese_dict,
                sheet_name="空调最大负荷预测数据",
                file_name="max_p_kt_predict_data",
            )
        return cls.return_common_func(result)
